#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HF 模型下载 GUI 工具。自动使用 hf-mirror.com 镜像，实时显示命令行输出
采用调用 hf download 命令方式，需确保已正确安装 huggingface_hub：
    pip install huggingface_hub
    然后在终端运行 hf download 命令，确保可执行。

repo_id 示例：speechbrain/spkrec-ecapa-voxceleb

官方推荐方法：
# Windows CMD
set HF_ENDPOINT=https://hf-mirror.com  
hf download ^
  --repo-type model ^
  speechbrain/spkrec-ecapa-voxceleb ^
  --local-dir models/speechbrain/spkrec-ecapa-voxceleb ^
"""
import os
import sys
import threading
import tkinter as tk
from tkinter import ttk, scrolledtext
from pathlib import Path
import subprocess


# ==============================
# 1. 环境配置
# ==============================
HF_MIRROR_URL = "https://hf-mirror.com"
os.environ["HF_ENDPOINT"] = HF_MIRROR_URL

# 获取系统 HF_TOKEN（如果存在）
HF_TOKEN = os.environ.get("HF_TOKEN", "").strip() or None


# ==============================
# 2. 下载函数（调用 hf download）
# ==============================
def download_model(repo_id, repo_type, token, log_widget, on_complete_callback):
    def safe_log(msg):
        log_widget.after(0, lambda: log_widget.insert(tk.END, msg + "\n"))
        log_widget.after(0, lambda: log_widget.see(tk.END))

    if not repo_id.strip():
        safe_log("❌ 错误：repo_id 不能为空！")
        on_complete_callback()
        return

    safe_path = repo_id.replace("/", os.sep)
    script_dir = Path(__file__).parent.resolve()
    local_dir = script_dir / "models" / safe_path
    local_dir.mkdir(parents=True, exist_ok=True)

    safe_log(f"📥 开始下载: {repo_id}")
    safe_log(f"📁 本地路径: {local_dir}")

    # 尝试找到 hf 命令
    hf_cmd = None
    # 方法1: 直接使用 'hf'（期望在 PATH 中）
    try:
        result = subprocess.run(["hf", "--help"], capture_output=True, text=True, timeout=5)
        if result.returncode == 0:
            hf_cmd = ["hf"]
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    # 方法2: 使用 python -m huggingface_hub.cli.hf（新版支持）
    if hf_cmd is None:
        try:
            result = subprocess.run([sys.executable, "-m", "huggingface_hub.cli.hf", "--help"],
                                    capture_output=True, text=True, timeout=5)
            if result.returncode == 0:
                hf_cmd = [sys.executable, "-m", "huggingface_hub.cli.hf"]
        except (subprocess.TimeoutExpired, Exception):
            pass

    if hf_cmd is None:
        safe_log("❌ 错误：未找到可用的 'hf' 命令。请升级 huggingface_hub：pip install -U huggingface_hub")
        on_complete_callback()
        return

    # 构建完整命令
    cmd = hf_cmd + [
        "download",
        "--repo-type", repo_type,
        "--local-dir", str(local_dir),
        repo_id
    ]

    if token:
        cmd.extend(["--token", token])

    safe_log(f"🔧 执行命令: {' '.join(cmd)}")

    env = os.environ.copy()
    env["HF_ENDPOINT"] = "https://hf-mirror.com"

    try:
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            bufsize=1,
            universal_newlines=True,
            env=env
        )

        for line in iter(process.stdout.readline, ''):
            if line:
                safe_log(line.rstrip())

        process.wait()

        if process.returncode == 0:
            safe_log("✅ 下载成功！")
        else:
            safe_log("❌ 下载失败（命令返回非零状态）")

    except Exception as e:
        safe_log(f"❌ 执行异常: {e}")
    finally:
        on_complete_callback()

# ==============================
# 3. GUI 主窗口
# ==============================
class HFDownloadGUI:
    def __init__(self, root):
        self.root = root
        root.title("HF 模型下载工具 (hf download + hf-mirror.com)")
        root.geometry("650x550")
        root.resizable(True, True)

        # Repo ID 输入
        ttk.Label(root, text="Repo ID:").grid(row=0, column=0, padx=10, pady=10, sticky="w")
        self.repo_id_var = tk.StringVar()
        ttk.Entry(root, textvariable=self.repo_id_var, width=60).grid(row=0, column=1, padx=5, pady=10, sticky="ew")

        # 类型下拉框
        ttk.Label(root, text="类型:").grid(row=1, column=0, padx=10, pady=5, sticky="w")
        self.repo_type_var = tk.StringVar(value="model")
        repo_types = ["model", "dataset", "space"]
        ttk.Combobox(root, textvariable=self.repo_type_var, values=repo_types, state="readonly", width=20).grid(
            row=1, column=1, padx=5, pady=5, sticky="w"
        )

        # Token 输入（可选）
        ttk.Label(root, text="Token (可选):").grid(row=2, column=0, padx=10, pady=5, sticky="w")
        self.token_var = tk.StringVar()
        ttk.Entry(root, textvariable=self.token_var, width=60).grid(row=2, column=1, padx=5, pady=5, sticky="ew")

        # 下载按钮
        self.download_btn = ttk.Button(root, text="开始下载", command=self.start_download)
        self.download_btn.grid(row=3, column=0, columnspan=2, pady=15)

        # 日志输出框
        ttk.Label(root, text="日志:").grid(row=4, column=0, padx=10, sticky="nw")
        self.log_text = scrolledtext.ScrolledText(root, height=18, state="normal", wrap=tk.WORD)
        self.log_text.grid(row=5, column=0, columnspan=2, padx=10, pady=5, sticky="nsew")

        root.grid_rowconfigure(5, weight=1)
        root.grid_columnconfigure(1, weight=1)

    def start_download(self):
        repo_id = self.repo_id_var.get().strip()
        if not repo_id:
            self.log_text.insert(tk.END, "⚠️ 请输入有效的 Repo ID！\n")
            return

        token = (self.token_var.get().strip() or HF_TOKEN) or None
        repo_type = self.repo_type_var.get()

        self.download_btn.config(state="disabled")
        self.log_text.delete(1.0, tk.END)

        def on_complete():
            self.root.after(0, lambda: self.download_btn.config(state="normal"))

        thread = threading.Thread(
            target=download_model,
            args=(repo_id, repo_type, token, self.log_text, on_complete),
            daemon=True
        )
        thread.start()


# ==============================
# 4. 入口
# ==============================
if __name__ == "__main__":
    root = tk.Tk()
    app = HFDownloadGUI(root)
    root.mainloop()